package com.example.mybatis.datasource;

import com.example.common.utils.Statistics;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.annotation.Primary;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import org.springframework.stereotype.Component;

import javax.sql.DataSource;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author zty
 * @apiNote
 * @project_name mybatis-demo
 * @user tain
 * @create_at 2023/6/23 9:02
 * @create_vio IntelliJ IDEA
 * 将RoutingDataSource替换DataSource后Druid不生效，走的是Hikari连接池.
 */
@Slf4j
@Primary
@Component
public class MyRoutingDataSource extends AbstractRoutingDataSource {

    private final Set<String> myLookupKeys = Collections.newSetFromMap(new ConcurrentHashMap<>());

    private final ThreadLocal<String> lookupKey = ThreadLocal.withInitial(() -> null);

    private MyRoutingDataSource(){
    }

    @Override
    public void afterPropertiesSet() {
        String dataSourceName1 = "myDataSource1";
        DataSource dataSource1 = MyDataSourceStore.getFirst();
        String dataSourceName2 = "myDataSource2";
        DataSource dataSource2 =MyDataSourceStore.getSecond();
        Map<Object, Object> targetDataSources = new HashMap<>();
        myLookupKeys.add(dataSourceName1);
        myLookupKeys.add(dataSourceName2);
        targetDataSources.put(dataSourceName1, dataSource1);
        targetDataSources.put(dataSourceName2, dataSource2);
        setTargetDataSources(targetDataSources);
        setDefaultTargetDataSource(dataSource1);
        super.afterPropertiesSet();
    }

    @Override
    protected String determineCurrentLookupKey() {
        Statistics.logInfo(log, "determineCurrentLookupKey", lookupKey.get());
        return lookupKey.get();
    }

    public Set<String> getMyLookupKeys() {
        return myLookupKeys;
    }

    public void reset(){
        lookupKey.remove();
    }

    public void setLookupKey(String lookupKey){
        this.lookupKey.set(lookupKey);
    }
}
